import torch
import torch.nn as nn


class Vgg16_net(nn.Module):
    def __init__(self, num_classes=10):
        super(Vgg16_net, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.flatten = nn.Flatten()

        self.classifier = nn.Sequential(
            nn.Linear(14 * 14 * 256, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(1024, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.flatten(x)
        x = self.classifier(x)
        return x


# 测试这个代码
# 假设10个分类
# num_classes = 10
# model = MyCNN(num_classes)
#
# # Example input image
# input_image = torch.randn(1, 3, 112, 112)  # Batch size 1, 3 channels, 112x112 image
# output = model(input_image)
# print(output.tolist())
# print(torch.argmax(output))
